from typing import Dict, List, Optional
from instruction_re.core.datatypes import Instance, Preffix, Relation
from ..instances.Answer import AnswerFormatter
from ..instances.EntityFormatter import Language
from ..instances import InstanceFormatter


class QuestionFormatter:
    available_languages = [Language.EN.value]
    pattern = "What does {0} {1}?"

    @classmethod
    def from_relations(
        cls, relations: Optional[List[Dict[str, List[str]]]], language: str = "en"
    ) -> Optional[str]:

        instrances = []
        for r in relations:
            instrances.append(cls.from_relation_triple(r, language))
        answer = ", ".join(instrances) + "."
        return answer

    @classmethod
    def from_relation_triple(
        cls, relation: Optional[Dict[str, List[str]]], language: str = "en"
    ) -> Optional[str]:
        if language not in cls.available_languages:
            raise ValueError(
                f"Expected language to be one of {cls.available_languages}"
            )

        if relation is None:
            return None

        return cls.pattern.format(relation["head"]["name"], relation["type"].lower())


class ObjExtractFormatter(InstanceFormatter):
    """
    Task: Given sentence, subject and the relation type extract the object
    """

    def format_instance(
        self,
        context: str,
        relations: Optional[List[Dict[str, List[str]]]],
        options: Optional[List[str]],
        instruction: str,
        examples: List[Relation] = None,
    ) -> Instance:

        instruction = Preffix.INSTRUCTION.value + instruction

        sentence = Preffix.CONTEXT.value + context

        # answer_format = Preffix.ANSWER.value + AnswerFormatter.from_relations(relations)
        answer_format = AnswerFormatter.from_relations(relations)

        # entity = EntityFormatter.from_relations(relations)
        relation_triples = [Relation.from_json(r) for r in relations]
        question = Preffix.QUESTION.value + QuestionFormatter.from_relations(relations)

        instance = Instance(
            instruction=instruction,
            context=sentence + " " + question + " " + Preffix.ANSWER.value,
            answer=answer_format,
            relations=relation_triples,
        )

        return instance


def main():
    sample = {
        "sentence": "John Wilkes Booth , who assassinated President Lincoln.",
        "relations": [
            {
                "head": {
                    "name": "John Wilkes Booth",
                    "type": "People",
                    "pos": [0, 17],
                },
                "type": "Kill",
                "tail": {
                    "name": "President Lincoln",
                    "type": "People",
                    "pos": [37, 54],
                },
            }
        ],
        "instruction": "Please extract all object from according to input sentence, subject and relation type",
    }

    formatter = ObjExtractFormatter()
    print(
        formatter.format_instance(
            context=sample["sentence"],
            relations=sample["relations"],
            instruction=sample["instruction"],
            options=[
                "Organization based in",
                "Located in",
                "Live in",
                "Work for",
                "Kill",
            ],
        )
    )


if __name__ == "__main__":
    main()
